package protocol

import (
	"container/list"
	"fmt"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"
)

var (
	NETFILTER_DROP   = NETFILTER_ACTION(1)
	NETFILTER_ACCEPT = NETFILTER_ACTION(2)
)

const fragentBufferSize = 65536
const maxFragmentSize = 1500 // 标准以太网MTU
const ipHeaderSize = 20
const ethernetHeaderSize = 14 // 以太网头部大小
const (
	FragmentLRUSize = 10000
	FragmentTimeout = 3 * time.Second
)

type NETFILTER_ACTION int

// todo 直接用xp保存分片内存
func FragmentPayload(ipv4Hdr header.IPv4, mtu int) (fragments [][]byte) {
	ipHeaderLength := ipv4Hdr.HeaderLength()
	baseMaxData := mtu - ethernetHeaderSize - int(ipHeaderLength) // 例如：1466
	maxDataPerFragment := (baseMaxData / 8) * 8
	if len(ipv4Hdr.Payload()) <= maxDataPerFragment {
		fragments = [][]byte{ipv4Hdr}
		return
	}
	originalPayload := ipv4Hdr.Payload()
	payloadLen := len(originalPayload)
	fragmentID := ipv4Hdr.ID()
	// 计算需要多少个分片
	numFragments := (payloadLen + maxDataPerFragment - 1) / maxDataPerFragment
	fragments = make([][]byte, numFragments)
	for i := 0; i < numFragments; i++ {
		// 计算当前分片的数据
		offset := i * maxDataPerFragment

		remainingData := payloadLen - offset
		currentFragmentDataSize := maxDataPerFragment
		if remainingData < maxDataPerFragment {
			currentFragmentDataSize = remainingData
		}

		// 创建分片
		fragmentPacket := make([]byte, int(ipHeaderLength)+currentFragmentDataSize)
		fragmentIPHdr := header.IPv4(fragmentPacket)

		// 复制IP头部
		copy(fragmentPacket[:ipHeaderLength], ipv4Hdr[:ipHeaderLength])

		// 复制分片数据
		copy(fragmentPacket[ipHeaderLength:], originalPayload[offset:offset+currentFragmentDataSize])

		// 设置分片相关字段
		fragmentIPHdr.SetTotalLength(uint16(len(fragmentPacket)))
		fragmentIPHdr.SetID(fragmentID) // 使用相同的ID

		// 设置分片偏移和标志
		var flags uint8

		if i < numFragments-1 {
			// 不是最后一个分片，设置MF标志
			flags = 1 // More Fragments flag
		} else {
			// 最后一个分片，不设置MF标志
			flags = 0 // No More Fragments flag
		}

		// 组合flags和fragment offset
		fragmentIPHdr.SetFlagsFragmentOffset(flags, uint16(offset))

		// 重新计算IP校验和
		fragmentIPHdr.SetChecksum(0)                  // 清零校验和字段
		checksum := fragmentIPHdr.CalculateChecksum() // 计算新校验和
		fragmentIPHdr.SetChecksum(^checksum)
		fragments[i] = fragmentIPHdr
	}
	return
}

// FragmentPayloadIpv6 对IPv6包进行分片，返回每个分片的完整数据（含IPv6头和Fragment头）
func FragmentPayloadIpv6(ipv6Pkg []byte, mtu int) ([][]byte, error) {
	const ipv6HeaderSize = 40
	const fragmentHeaderSize = 8

	if len(ipv6Pkg) < ipv6HeaderSize {
		return nil, fmt.Errorf("packet too short for IPv6 header")
	}
	if header.IPVersion(ipv6Pkg) != header.IPv6Version {
		return nil, fmt.Errorf("not an IPv6 packet")
	}

	ipv6Hdr := header.IPv6(ipv6Pkg)
	payload := ipv6Hdr.Payload()
	payloadLen := len(payload)

	// 分片数据最大长度，需8字节对齐
	maxDataPerFragment := ((mtu - ethernetHeaderSize - ipv6HeaderSize - fragmentHeaderSize) / 8) * 8
	if maxDataPerFragment <= 0 {
		return nil, fmt.Errorf("mtu too small for fragmentation")
	}
	if payloadLen <= maxDataPerFragment {
		// 不需要分片
		return [][]byte{ipv6Pkg}, nil
	}

	// RFC要求分片ID为32位
	fragmentID := uint32(time.Now().UnixNano() & 0xffffffff)
	nextHeader := ipv6Hdr.NextHeader()

	numFragments := (payloadLen + maxDataPerFragment - 1) / maxDataPerFragment
	fragments := make([][]byte, numFragments)

	for i := 0; i < numFragments; i++ {
		offset := i * maxDataPerFragment
		remaining := payloadLen - offset
		dataLen := maxDataPerFragment
		if remaining < maxDataPerFragment {
			dataLen = remaining
		}

		// 构造分片包：IPv6头 + Fragment头 + 分片数据
		fragPacket := make([]byte, ipv6HeaderSize+fragmentHeaderSize+dataLen)

		// 复制IPv6头
		copy(fragPacket[:ipv6HeaderSize], ipv6Pkg[:ipv6HeaderSize])
		fragIPv6 := header.IPv6(fragPacket)
		fragIPv6.SetPayloadLength(uint16(fragmentHeaderSize + dataLen))
		fragIPv6.SetNextHeader(header.IPv6FragmentHeader)

		// 构造Fragment Header
		fragHdr := header.IPv6Fragment(fragPacket[ipv6HeaderSize : ipv6HeaderSize+fragmentHeaderSize])
		SetNextHeader(fragHdr, nextHeader) // 设置原始的Next Header, tcp udp icmp等
		SetIdentification(fragHdr, fragmentID)
		SetFragmentOffset(fragHdr, uint16(offset/8))
		if i < numFragments-1 {
			SetMFlag(fragHdr)
		} else {
			ClearMFlag(fragHdr)
		}

		// 复制分片数据
		copy(fragPacket[ipv6HeaderSize+fragmentHeaderSize:], payload[offset:offset+dataLen])

		fragments[i] = fragPacket
	}

	return fragments, nil
}

// SetNextHeader 设置 Fragment Header 的 Next Header 字段
func SetNextHeader(f header.IPv6Fragment, nextHeader uint8) {
	f[0] = nextHeader
}

// SetFragmentOffset 设置 Fragment Offset 字段（单位8字节）
func SetFragmentOffset(f header.IPv6Fragment, offset uint16) {
	// offset: 13 bits, bits 3-15 of bytes 2-3
	// 清除原有 offset
	f[2] &^= 0xF8
	f[3] = 0
	// 设置 offset
	f[2] |= byte((offset >> 5) & 0xF8)
	f[3] |= byte((offset << 3) & 0xF8)
}

// SetMFlag 设置 More Fragments 标志
func SetMFlag(f header.IPv6Fragment) {
	f[2] |= 0x01
}

// ClearMFlag 清除 More Fragments 标志
func ClearMFlag(f header.IPv6Fragment) {
	f[2] &^= 0x01
}

// SetIdentification 设置 Identification 字段
func SetIdentification(f header.IPv6Fragment, id uint32) {
	f[4] = byte(id >> 24)
	f[5] = byte(id >> 16)
	f[6] = byte(id >> 8)
	f[7] = byte(id)
}

func GetPortFromIpv4FristFragment(ipv4Hdr header.IPv4) (srcPort, dstPort uint16, protocol uint8, err error) {
	nextHeader := ipv4Hdr.Protocol()
	switch tcpip.TransportProtocolNumber(nextHeader) {
	case header.TCPProtocolNumber:
		if len(ipv4Hdr.Payload()) < 20 {
			err = fmt.Errorf("tcp header too short")
			return
		}
		tcpHdr := header.TCP(ipv4Hdr.Payload())
		srcPort = tcpHdr.SourcePort()
		dstPort = tcpHdr.DestinationPort()
		protocol = uint8(header.TCPProtocolNumber)
	case header.UDPProtocolNumber:
		if len(ipv4Hdr.Payload()) < 8 {
			err = fmt.Errorf("udp header too short")
			return
		}
		udpHdr := header.UDP(ipv4Hdr.Payload())
		srcPort = udpHdr.SourcePort()
		dstPort = udpHdr.DestinationPort()
		protocol = uint8(header.UDPProtocolNumber)
	case header.ICMPv4ProtocolNumber:
		icmp := header.ICMPv4(ipv4Hdr.Payload())
		srcPort = icmp.Ident()
		dstPort = icmp.Ident()
		protocol = uint8(header.ICMPv4ProtocolNumber)
	default:
		err = fmt.Errorf("unsupported protocol: %d", ipv4Hdr.Protocol())
	}
	return
}

// IPv4分片信息
type IPv4Fragment struct {
	ID          uint16        // 标识符
	SrcIP       tcpip.Address // 源IP
	DstIP       tcpip.Address // 目标IP
	Protocol    uint8         // 协议
	Fragments   *list.List    // 分片列表
	TotalLength uint16        // 总长度（当收到最后一个分片时设置）
	ReceivedLen uint16        // 已接收长度
	LastUpdate  time.Time     // 最后更新时间
	Complete    bool          // 是否完整
}

// 单个分片
type FragmentPiece struct {
	Offset uint16 // 偏移量（以8字节为单位）
	Length uint16 // 数据长度
	Data   []byte // 分片数据
	MF     bool   // More Fragments标志
}

// IPv4分片重组器
type IPv4Reassembler struct {
	fragments map[string]*IPv4Fragment // key: srcIP_dstIP_protocol_id
	mutex     sync.Mutex
	maxAge    time.Duration // 分片最大存活时间
}

var DefaultIPv4Reassembler = NewIPv4Reassembler()

func NewIPv4Reassembler() *IPv4Reassembler {
	r := &IPv4Reassembler{
		fragments: make(map[string]*IPv4Fragment),
		maxAge:    30 * time.Second, // 30秒超时
	}

	// 启动清理goroutine
	go r.cleanup()
	return r
}

// 生成分片键
func (r *IPv4Reassembler) getFragmentKey(srcIP, dstIP tcpip.Address, protocol uint8, id uint16) string {
	return fmt.Sprintf("%s_%s_%d_%d", srcIP.String(), dstIP.String(), protocol, id)
}

// 检查是否为分片包
func IsIPv4Fragment(ipv4Hdr header.IPv4) bool {
	flags := ipv4Hdr.Flags()
	fragOffset := ipv4Hdr.FragmentOffset()

	// 如果设置了MF标志或者偏移量不为0，则为分片包
	return (flags&header.IPv4FlagMoreFragments) != 0 || fragOffset != 0
}

// 添加分片
func (r *IPv4Reassembler) AddFragment(ipv4Hdr header.IPv4) (complete bool, reassembled header.IPv4, err error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	srcIP := ipv4Hdr.SourceAddress()
	dstIP := ipv4Hdr.DestinationAddress()
	protocol := ipv4Hdr.Protocol()
	id := ipv4Hdr.ID()
	fragOffset := ipv4Hdr.FragmentOffset()
	flags := ipv4Hdr.Flags()
	mf := (flags & header.IPv4FlagMoreFragments) != 0

	key := r.getFragmentKey(srcIP, dstIP, protocol, id)

	// 获取或创建分片信息
	fragInfo, exists := r.fragments[key]
	if !exists {
		fragInfo = &IPv4Fragment{
			ID:         id,
			SrcIP:      srcIP,
			DstIP:      dstIP,
			Protocol:   protocol,
			Fragments:  list.New(),
			LastUpdate: time.Now(),
		}
		r.fragments[key] = fragInfo
	}

	fragInfo.LastUpdate = time.Now()

	// 创建分片片段
	payload := ipv4Hdr.Payload()
	piece := &FragmentPiece{
		Offset: fragOffset,
		Length: uint16(len(payload)),
		Data:   make([]byte, len(payload)),
		MF:     mf,
	}
	copy(piece.Data, payload)

	// 如果是最后一个分片，设置总长度
	if !mf {
		fragInfo.TotalLength = fragOffset + piece.Length
	}

	// 插入分片（按偏移量排序）
	inserted := false
	for e := fragInfo.Fragments.Front(); e != nil; e = e.Next() {
		existing := e.Value.(*FragmentPiece)
		if piece.Offset < existing.Offset {
			fragInfo.Fragments.InsertBefore(piece, e)
			inserted = true
			break
		} else if piece.Offset == existing.Offset {
			// 重复分片，替换
			e.Value = piece
			inserted = true
			break
		}
	}

	if !inserted {
		fragInfo.Fragments.PushBack(piece)
	}

	// 更新已接收长度
	fragInfo.ReceivedLen += piece.Length

	// 检查是否完整
	if fragInfo.TotalLength > 0 {
		complete = r.checkComplete(fragInfo)
		if complete {
			reassembled, err = r.reassemble(fragInfo, ipv4Hdr)
			if err == nil {
				delete(r.fragments, key) // 删除已完成的分片
			}
		}
	}

	return complete, reassembled, err
}

// 检查分片是否完整
func (r *IPv4Reassembler) checkComplete(fragInfo *IPv4Fragment) bool {
	if fragInfo.TotalLength == 0 {
		return false
	}

	expectedLen := uint16(0)
	currentOffset := uint16(0)

	for e := fragInfo.Fragments.Front(); e != nil; e = e.Next() {
		piece := e.Value.(*FragmentPiece)

		// 检查是否有空隙
		if piece.Offset != currentOffset {
			return false
		}

		expectedLen += piece.Length
		currentOffset = piece.Offset + piece.Length
	}

	return expectedLen == fragInfo.TotalLength
}

// 重组分片
func (r *IPv4Reassembler) reassemble(fragInfo *IPv4Fragment, originalHdr header.IPv4) (header.IPv4, error) {
	// 计算新的IP包总长度
	newTotalLen := header.IPv4MinimumSize + int(fragInfo.TotalLength)

	// 创建新的IP包
	newPacket := make([]byte, newTotalLen)

	// 复制原始IP头部
	copy(newPacket[:header.IPv4MinimumSize], originalHdr[:header.IPv4MinimumSize])

	// 创建新的IPv4头部
	newHdr := header.IPv4(newPacket)

	// 设置新的长度
	newHdr.SetTotalLength(uint16(newTotalLen))

	// 清除分片相关字段
	newHdr.SetFlagsFragmentOffset(0, 0)

	// 重组payload
	payloadOffset := header.IPv4MinimumSize
	for e := fragInfo.Fragments.Front(); e != nil; e = e.Next() {
		piece := e.Value.(*FragmentPiece)
		copy(newPacket[payloadOffset:], piece.Data)
		payloadOffset += int(piece.Length)
	}

	// 重新计算IP头部校验和
	newHdr.SetChecksum(0)
	checksum := CalculateChecksum(newHdr[:newHdr.HeaderLength()], 0)
	newHdr.SetChecksum(checksum)

	return newHdr, nil
}

// 清理过期分片
func (r *IPv4Reassembler) cleanup() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		r.mutex.Lock()
		now := time.Now()
		for key, fragInfo := range r.fragments {
			if now.Sub(fragInfo.LastUpdate) > r.maxAge {
				delete(r.fragments, key)
			}
		}
		r.mutex.Unlock()
	}
}

type FragmentKey struct {
	SrcIP    tcpip.Address
	DstIP    tcpip.Address
	Protocol uint8
	ID       uint32
}

// IpFragmentInfo 分片信息
type IpFragmentInfo struct {
	lastSeen     time.Time
	originPort   uint16                        // UDP/TCP端口或ICMPv6 ID
	fragments    [][]byte                      // 缓存分片数据
	hasL4Headers bool                          // 标记是否已找到包含L4头信息的分片
	protocol     tcpip.TransportProtocolNumber // 协议类型 (UDP/TCP/ICMPv6)
}

type IpFragmentManager struct {
	mu    sync.Mutex
	cache *lru.Cache[FragmentKey, *IpFragmentInfo]
}

var DefaultIpFragmentManager = NewIpFragmentManager()

func NewIpFragmentManager() *IpFragmentManager {
	cache, err := lru.New[FragmentKey, *IpFragmentInfo](FragmentLRUSize)
	if err != nil {
		panic(err)
	}

	manager := &IpFragmentManager{
		cache: cache,
	}

	// 启动定期清理协程
	go manager.startCleanupRoutine()

	return manager
}

// ProcessFragment 处理IPv6分片
// 返回: (是否应该立即处理, 原始端口, 所有分片数据)
func (m *IpFragmentManager) ProcessIpFragment(ipVersion uint8, ipPkg []byte) (shouldProcess bool, originPort uint16, allFragments [][]byte) {
	key := FragmentKey{}
	var hasL4Info bool
	var l4Data []byte
	var port uint16
	var nextHeader tcpip.TransportProtocolNumber
	switch ipVersion {
	case header.IPv4Version:
		ipv4Hdr := header.IPv4(ipPkg)
		flags := ipv4Hdr.Flags()
		fragOffset := ipv4Hdr.FragmentOffset()
		// first fragment
		if (flags&header.IPv4FlagMoreFragments) != 0 || fragOffset == 0 {
			hasL4Info = true
		}
		l4Data = ipv4Hdr.Payload()
		nextHeader = tcpip.TransportProtocolNumber(ipv4Hdr.Protocol())
	case header.IPv6Version:
		ipv6Hdr := header.IPv6(ipPkg)
		payload := ipv6Hdr.Payload()
		if len(payload) < header.IPv6FragmentHeaderSize {
			return false, 0, nil
		}
		fragHdr := header.IPv6Fragment(payload)
		if !fragHdr.IsValid() {
			return false, 0, nil
		}
		key.SrcIP, key.DstIP, key.ID = ipv6Hdr.SourceAddress(), ipv6Hdr.DestinationAddress(), fragHdr.ID()
		nextHeader = tcpip.TransportProtocolNumber(fragHdr.NextHeader())
		// 检查是否包含L4头信息（offset=0且是支持的协议）
		hasL4Info = (nextHeader == header.UDPProtocolNumber ||
			nextHeader == header.TCPProtocolNumber ||
			nextHeader == header.ICMPv6ProtocolNumber) && fragHdr.FragmentOffset() == 0 && fragHdr.More()
		l4Data = payload[header.IPv6FragmentHeaderSize:]
	default:
		panic("unknown ip version")
	}

	if hasL4Info {
		port = m.extractPortFromFragment(l4Data, nextHeader)
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	// 1. 如果分片包含port或者icmp id
	if hasL4Info {
		fragInfo := &IpFragmentInfo{
			lastSeen:     time.Now(),
			originPort:   port,
			hasL4Headers: true,
			protocol:     nextHeader,
		}
		// 检查是否已有缓存的分片
		allFragments = append(allFragments, [][]byte{ipPkg}...)
		if existingInfo, exists := m.cache.Get(key); exists {
			// 合并已缓存的分片
			allFragments = append(allFragments, existingInfo.fragments...)
		}

		// 直接更新lru，并设置originPort
		m.cache.Add(key, fragInfo)
		return true, port, allFragments
	}
	// 2. 查询LRU缓存
	if existingInfo, exists := m.cache.Get(key); exists {
		now := time.Now()
		// 2.1 if lastSeen >= 3s
		if now.Sub(existingInfo.lastSeen) >= FragmentTimeout {
			// 删除旧结构并重新保存，hasL4Headers=false
			newInfo := &IpFragmentInfo{
				lastSeen:     now,
				originPort:   0,
				fragments:    [][]byte{ipPkg},
				hasL4Headers: false,
				protocol:     nextHeader,
			}
			m.cache.Add(key, newInfo)
			return false, 0, nil
		}
		// 2.2 if lastSeen < 3s
		if existingInfo.hasL4Headers {
			// 已经有L4头信息，直接返回
			allFragments = append(allFragments, [][]byte{ipPkg}...)
			return true, existingInfo.originPort, allFragments
		}
		// 继续缓存分片
		existingInfo.lastSeen = now
		existingInfo.fragments = append(existingInfo.fragments, ipPkg)
		m.cache.Add(key, existingInfo) // 更新LRU位置
		return false, 0, nil
	}

	// 3. 查询失败，当前分片不包含l4头信息，且LRU中没有缓存
	newInfo := &IpFragmentInfo{
		lastSeen:     time.Now(),
		originPort:   0,
		fragments:    [][]byte{ipPkg},
		hasL4Headers: false,
		protocol:     nextHeader,
	}
	m.cache.Add(key, newInfo)
	return false, 0, nil
}

// extractPortFromFragment 从分片数据中提取端口或ICMPv6 ID
func (m *IpFragmentManager) extractPortFromFragment(l4Data []byte, protocol tcpip.TransportProtocolNumber) uint16 {
	switch protocol {
	case header.UDPProtocolNumber:
		udp := header.UDP(l4Data)
		return udp.DestinationPort()
	case header.TCPProtocolNumber:
		tcp := header.TCP(l4Data)
		return tcp.DestinationPort()
	case header.ICMPv6ProtocolNumber:
		icmpv6 := header.ICMPv6(l4Data)
		return icmpv6.Ident()
	}
	return 0
}

// startCleanupRoutine 启动定期清理协程
func (m *IpFragmentManager) startCleanupRoutine() {
	ticker := time.NewTicker(1 * time.Second) // 每秒检查一次
	defer ticker.Stop()

	for range ticker.C {
		m.cleanupExpired()
	}
}

// cleanupExpired 清理过期的分片
func (m *IpFragmentManager) cleanupExpired() {
	now := time.Now()

	// 获取所有键
	keys := m.cache.Keys()

	for _, key := range keys {
		if fragInfo, exists := m.cache.Peek(key); exists { // 使用Peek避免影响LRU顺序
			if now.Sub(fragInfo.lastSeen) >= FragmentTimeout {
				m.cache.Remove(key)
			}
		}
	}
}
